'''
Description: 
Version: 2.0
Date: 2022-05-13 17:45:44
LastEditors: Meroke
<<<<<<< HEAD
LastEditTime: 2022-05-18 16:56:53
=======
LastEditTime: 2022-07-22 17:11:34
>>>>>>> origin/develop
'''
from hashlib import new
import cv2
from matplotlib import image
import numpy as np
import sys
import os

from numpy import vstack


image_size = 1024

def nothing(x):
    pass



'''
description:  HSV 阈值调节可视化
param path 图片路径
return {*}
'''
def hsv_adjust(path):
    cv2.namedWindow('mask')
    # 创建Trackbar
    cv2.createTrackbar('LR', 'mask', 0, 255, nothing)
    cv2.createTrackbar('LG', 'mask', 0, 255, nothing)
    cv2.createTrackbar('LB', 'mask', 0, 255, nothing)
    cv2.createTrackbar('HR', 'mask', 0, 255, nothing)
    cv2.createTrackbar('HG', 'mask', 0, 255, nothing)
    cv2.createTrackbar('HB', 'mask', 0, 255, nothing)

    img = cv2.imread(path)
    img = cv2.resize(img,(512,512))
    hsv_img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
    # 设置初值
    cv2.setTrackbarPos('LR','mask',98)
    cv2.setTrackbarPos('LG','mask',78)
    cv2.setTrackbarPos('LB','mask',39)
    cv2.setTrackbarPos('HR','mask',125)
    cv2.setTrackbarPos('HG','mask',255)
    cv2.setTrackbarPos('HB','mask',255)
    show("origin",img)
    cv2.moveWindow("origin",512,512)
    while(1): # RGB -> SHV 对应变换
        # 循环更新
        LR = cv2.getTrackbarPos('LR','mask')
        LG = cv2.getTrackbarPos('LG','mask')
        LB = cv2.getTrackbarPos('LB','mask')
        HR = cv2.getTrackbarPos('HR','mask')
        HG = cv2.getTrackbarPos('HG','mask')
        HB = cv2.getTrackbarPos('HB','mask')
        lower = np.array([LR, LG, LB])  # 所要检测的像素范围
        upper = np.array([HR, HG, HB])  # 此处检测黄色区域
        mask = cv2.inRange(hsv_img, lowerb=lower, upperb=upper)
        
        cv2.imshow("mask", mask)
        k = cv2.waitKey(1) & 0xFF
        if k == 27:
            break
    cv2.destroyAllWindows()

# 计算轮廓面积，被调用排序轮廓
def cnt_area(cnt):
  area = cv2.contourArea(cnt)
  return area

"""
    
    image's box_list index:
    0    1
    3    2
"""
'''
description: 轮廓四角下标重编排，并扩大四角范围
param box 轮廓
param bigger_size  扩大范围
return {*}
'''
def bigger_box(box,bigger_size=10):
    points = [None]*4
    sorted_box = [None]*4
    # print(box)
    for i in range(len(box)):
        points[i] = box[i]
    points = sorted(points,key= lambda d: d[1],reverse=False)
    list12 = points[:2]
    list34 = points[2:4]
    # print("points",points)
    if list12[0][0] <= list12[1][0]:
        sorted_box[0]=list12[0]
        sorted_box[1]=list12[1]
    else:
        sorted_box[0]=list12[1]
        sorted_box[1]=list12[0]

    if list34[0][0] <= list34[1][0]:
        sorted_box[2]=list34[1]
        sorted_box[3]=list34[0] 
    else:
        sorted_box[2]=list34[0]
        sorted_box[3]=list34[1]  
    sorted_box[0][0] -= bigger_size
    sorted_box[0][1] -= bigger_size

    sorted_box[1][0] += bigger_size
    sorted_box[1][1] -= bigger_size

    sorted_box[2][0] += bigger_size
    sorted_box[2][1] += bigger_size

    sorted_box[3][0] -= bigger_size
    sorted_box[3][1] += bigger_size

    # print("list12",list12)
    # print("list34",list34)
    # print("sorted_box",sorted_box)
    for i in range(4):
        for j in range(2):
            if sorted_box[i][j] < 0:
                sorted_box[i][j] = 0
            elif(sorted_box[i][j]>image_size):
                sorted_box[i][j] = image_size

    return sorted_box

'''
description: HSV 分割测试例程
param path 图片路径
param name 图片名称
return {*}
'''
def hsv_split_show(path,name):
    img = cv2.imread(path)
    img = cv2.resize(img,(image_size,image_size))
    img = cv2.GaussianBlur(img,(3,3),0)
    hsv_img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
    lower = np.array([0, 0, 80])  # 所要检测的像素范围
    upper = np.array([70, 80, 255])  # 此处检测白色区域
    mask = cv2.inRange(hsv_img, lowerb=lower, upperb=upper)
    contours = cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)[0]
    contours.sort(key=cnt_area, reverse=True)  # 自大到小排序
    # print(contours[0])
    # cv2.drawContours(img,[contours[0]],-1,(0,0,255),1)
    new_mask = np.zeros((image_size,image_size),np.int8)
    for cnt in contours[:5]:
        rect = cv2.minAreaRect(cnt)
        box = cv2.boxPoints(rect)  # 获取最小外接矩形的4个顶点
        box = np.int0(box)
        box = bigger_box(box)
        box = np.array(box)
        # 限制书本类狭长物体被检测：
        l1 = box[1][0] - box[0][0]
        l2 = box[3][1] - box[0][1]
        print("l1",l1)
        print("l2",l2)
        cv2.putText(img, str(l1)+" "+str(l2) ,(box[0][0]+5,box[0][1]+25),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0), 1)
        if abs(l1/l2 -1) > 0.75 or \
            box[1][0] - box[0][0] < 150 or \
            box[3][1] - box[0][1]  < 150:
            continue
        cv2.polylines(img, [box], True, (0,0,255), 3)
        for i in range(len(box)):
            cv2.putText(img,str(i),(box[i][0]+2,box[i][1]+2),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0), 2)
        # mask = cv2.cvtColor(mask,cv2.COLOR_HSV2BGR)

        new_mask[box[0][1]:box[2][1], box[0][0]:box[2][0]] = 255
        # new_mask[box[1][1]:box[3][1], box[1][0]:box[3][0]] = 255

    cv2.imshow(name+"cnt",img)
    cv2.imshow(name, new_mask)
    masked_img = cv2.bitwise_and(img,img, mask=new_mask)
    cv2.imshow(name+"mask",masked_img)
    k = cv2.waitKey(0) & 0xFF
    if k == 27:
        sys.exit()
    cv2.destroyAllWindows()

'''
description:  图片展示函数
param name 图片名称
param img 图片
param mode 是否需要单独暂停显示，default 0,不开启
return {*}
'''
def show(name,img,mode=0):
    cv2.imshow(name,img)
    if(mode):
        cv2.waitKey(0)

'''
description:  hsv正式程序，分割蓝色区域
param img  传入图像
return {*}
'''
def hsv_split_pre(img,h,w):
    img = cv2.GaussianBlur(img,(3,3),0)
    hsv_img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
    lower = np.array([98, 78, 39])  # 所要检测的像素范围
    upper = np.array([125, 255, 255])  # 此处检测白色文字区域
    mask = cv2.inRange(hsv_img, lowerb=lower, upperb=upper)
    # show("mask",mask,1)
<<<<<<< HEAD
=======

    kernel_big = cv2.getStructuringElement(cv2.MORPH_RECT, (11,11))
    kernel_small = cv2.getStructuringElement(cv2.MORPH_RECT, (9,9))
    mask = cv2.erode(mask,kernel_big)
    # show("erode",mask,1)
    mask = cv2.dilate(mask,kernel_small)
    # show("dlite",mask,1)
    # cv2.imwrite("/home/tdb2/2022_jsjds/detection_2022/ocr/image/result/mask.jpg",mask)
>>>>>>> origin/develop
    contours = cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)[0]
    contours.sort(key=cnt_area, reverse=True)  # 自大到小排序
    # cv2.drawContours(img,[contours[0]],-1,(0,0,255),1)
    new_mask = np.zeros((h,w),np.int8)
    count = 0
    cnt_list = []
    for cnt in contours[:5]:
        rect = cv2.minAreaRect(cnt)
        box = cv2.boxPoints(rect)  # 获取最小外接矩形的4个顶点
        box = np.int0(box)
        box = bigger_box(box)
        box = np.array(box)
        # 限制书本类狭长物体被检测：
        l1 = box[1][0] - box[0][0]
        l2 = box[3][1] - box[0][1]
        short_len = l1 if l1 < l2 else l2
        if short_len<300 :
            continue
        count += 1
        new_mask[box[0][1]:box[2][1], box[0][0]:box[2][0]] = 255
        new_mask[box[1][1]:box[3][1], box[3][0]:box[1][0]] = 255
        cnt_list.append(box)
        # cv2.putText(img,str(l1)+" " + str(l2),(box[0][0]+20,box[0][1]+20),
        #             cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255), 1)
        # cv2.polylines(img, [box], True, (0,0,255), 2)
    if(count == 0):
        new_mask[:-1,:-1] = 255
    masked_img = cv2.bitwise_and(img,img, mask=new_mask)
<<<<<<< HEAD
    return masked_img
=======
    # show("masked",masked_img,1)
    # cv2.imwrite("/home/tdb2/2022_jsjds/detection_2022/ocr/image/result/masked_img.jpg",masked_img)
    print("cnt_centre_list",cnt_centre_list)
    # boxes_list.sort()
    boxes_list = sorted(boxes_list,key=lambda x:x[0][0])
    cnt_centre_list = sorted(cnt_centre_list,key=lambda x:x[0][0])
    extra_big_box = sorted(extra_big_box,key=lambda x:x[0])
    # cnt_centre_list.sort()
    # sorted(boxes_list)
    return masked_img,cnt_centre_list,boxes_list,extra_big_box
>>>>>>> origin/develop

'''
description: HSV 正式程序，分割白色
param img  传入图像
return {*} cnt_list,masked_img 轮廓列表  最终处理图像
'''
def hsv_split(img,h,w):
    # img = cv2.resize(img,(int(img.shape[0]/10),int(img.shape[1]/10)))
    img = cv2.GaussianBlur(img,(3,3),0)
    hsv_img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
    lower = np.array([0, 0, 80])  # 所要检测的像素范围
    upper = np.array([70, 80, 255])  # 此处检测白色文字区域
    mask = cv2.inRange(hsv_img, lowerb=lower, upperb=upper)
    # show("02-1mask",mask)
    contours = cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)[0]
    contours.sort(key=cnt_area, reverse=True)  # 自大到小排序
    # cv2.drawContours(img,[contours[0]],-1,(0,0,255),1)
    new_mask = np.zeros((h,w),np.int8)
    count = 0
    cnt_list = []
    for cnt in contours[:5]:
        rect = cv2.minAreaRect(cnt)
        box = cv2.boxPoints(rect)  # 获取最小外接矩形的4个顶点
        box = np.int0(box)
        box = bigger_box(box)
        box = np.array(box)
        # 限制书本类狭长物体被检测：
        l1 = box[1][0] - box[0][0]
        l2 = box[3][1] - box[0][1]
        short_len = l1 if l1 < l2 else l2
        # 筛选去长条状轮廓 和 尺寸过小、过大轮廓
        if abs(l1/l2 -1) > 0.75 or \
            short_len<150 or short_len > 512:
            continue
        count += 1
        # 分割白色区域
        new_mask[box[0][1]:box[2][1], box[0][0]:box[2][0]] = 255
        cnt_list.append(box)
        # 显示矩形框的长宽，显示对比检测效果
        # cv2.putText(img,str(l1)+" " + str(l2),(box[0][0]+20,box[0][1]+20),
        #             cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255), 1)
        # cv2.polylines(img, [box], True, (0,0,255), 2)
    # 未识别到轮廓，则返回原图
    if(count == 0):
        new_mask[:-1,:-1] = 255
    masked_img = cv2.bitwise_and(img,img, mask=new_mask)
    cnt_list = np.array(cnt_list)
    return cnt_list,masked_img

if __name__ == "__main__":
    source_image_path = "/home/meroke/code/detection/result/error/"
    # source_image_path = "/home/meroke/code/detection/train_images/"
    for i in os.listdir(source_image_path):
        if "jpg" in i:
            if "del" in i:
                continue
            print(i)
            hsv_adjust(source_image_path + i)
            # hsv_split_show(source_image_path + i,i)
            # img = cv2.imread(source_image_path + i)
            # img = cv2.resize(img,(image_size,image_size))
            # hsv_split_pre(img)
            print(i + " done!")